{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Following the calibrate_quant_resnet50 example, now we fine tune the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: Logging before flag parsing goes to stderr.\n",
      "W0608 21:25:39.018203 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
      "W0608 21:25:39.019082 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
      "W0608 21:25:39.019555 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
      "W0608 21:25:39.020030 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
      "W0608 21:25:39.020492 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
      "W0608 21:25:39.020947 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
      "W0608 21:25:39.021392 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n"
     ]
    }
   ],
   "source": [
    "import datetime\n",
    "import os\n",
    "import sys\n",
    "import time\n",
    "\n",
    "import torch\n",
    "import torch.utils.data\n",
    "from torch import nn\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "\n",
    "from pytorch_quantization import nn as quant_nn\n",
    "from pytorch_quantization import calib\n",
    "from pytorch_quantization.tensor_quant import QuantDescriptor\n",
    "\n",
    "from absl import logging\n",
    "logging.set_verbosity(logging.FATAL)  # Disable logging as they are too noisy in notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For simplicity, import train and eval functions from the train script from torchvision instead of copything them here\n",
    "sys.path.append(\"/raid/skyw/models/torchvision/references/classification/\")\n",
    "from train import evaluate, train_one_epoch, load_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "QuantResNet(\n",
       "  (conv1): QuantConv2d(\n",
       "    3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False\n",
       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.6387 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "  )\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (relu): ReLU(inplace=True)\n",
       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
       "  (layer1): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.9730 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.7266](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0971 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.4679](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3318 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3936](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.9730 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.9879](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4872 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.2618](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0466 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0564, 0.5201](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5106 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.2946](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5250 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0533, 0.1921](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9980 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0810, 0.2856](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.6532 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0197, 0.2752](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5499 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0670, 0.3532](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1606 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0694, 0.2993](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1425 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3917](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5499 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.5662](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4626 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0370, 0.2522](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8304 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0142, 0.2998](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1722 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3038](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4864 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0653, 0.2383](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9450 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0646, 0.2556](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8535 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0162, 0.3522](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5229 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0648, 0.2814](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9247 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0595, 0.2210](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9747 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0127, 0.2956](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5941 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0842, 0.3425](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3565 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0544, 0.2008](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0293 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3212](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5941 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3460](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3305 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0407, 0.2942](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9844 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0432, 0.2634](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9333 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.4969](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3388 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0469, 0.2715](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8617 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0397, 0.2100](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.7507 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3538](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (3): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3554 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0553, 0.2390](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9257 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0455, 0.2792](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8117 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3126](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (4): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4199 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0598, 0.2722](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9274 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0459, 0.1919](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8702 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0160, 0.3161](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (5): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4258 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0618, 0.3995](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.2256 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0498, 0.2236](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3560 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0180, 0.3288](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (layer4): Sequential(\n",
       "    (0): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3915 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0744, 0.3415](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1571 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0432, 0.3993](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1295 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0198, 0.3546](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "      (downsample): Sequential(\n",
       "        (0): QuantConv2d(\n",
       "          1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3915 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0197, 0.6413](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        )\n",
       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (1): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=3.9348 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0482, 0.7003](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1277 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0379, 0.2257](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8992 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0211, 0.2427](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "    (2): Bottleneck(\n",
       "      (conv1): QuantConv2d(\n",
       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.2181 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0601, 0.4541](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): QuantConv2d(\n",
       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.2051 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0338, 0.1416](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv3): QuantConv2d(\n",
       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1045 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0119, 0.2798](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "      )\n",
       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (relu): ReLU(inplace=True)\n",
       "    )\n",
       "  )\n",
       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
       "  (fc): QuantLinear(\n",
       "    in_features=2048, out_features=1000, bias=True\n",
       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.5345 calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1716, 0.7371](1000) calibrator=MaxCalibrator(track_amax=False) quant)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pytorch_quantization import quant_modules\n",
    "quant_modules.initialize()\n",
    "\n",
    "# Create and load the calibrated model\n",
    "model = torchvision.models.resnet50()\n",
    "model.load_state_dict(torch.load(\"/tmp/quant_resnet50-calibrated.pth\"))\n",
    "model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"/raid/data/imagenet/imagenet_pytorch\"\n",
    "\n",
    "traindir = os.path.join(data_path, 'train')\n",
    "valdir = os.path.join(data_path, 'val')\n",
    "dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, False, False)\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=256,\n",
    "    sampler=train_sampler, num_workers=4, pin_memory=True)\n",
    "\n",
    "data_loader_test = torch.utils.data.DataLoader(\n",
    "    dataset_test, batch_size=256,\n",
    "    sampler=test_sampler, num_workers=4, pin_memory=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantized fine tuning\n",
    "Let's fine tune the model with fake quantization. We only fine tune for 1 epoch as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)\n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)\n",
    "\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=128,\n",
    "    sampler=train_sampler, num_workers=16, pin_memory=True)\n",
    "\n",
    "# Training takes about one and half hour per epoch on single V100\n",
    "train_one_epoch(model, criterion, optimizer, data_loader, \"cuda\", 0, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate the fine tuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test:  [  0/196]  eta: 0:09:36  loss: 0.4680 (0.4680)  acc1: 85.9375 (85.9375)  acc5: 98.0469 (98.0469)  time: 2.9406  data: 2.1852  max mem: 16096\n",
      "Test:  [ 10/196]  eta: 0:01:58  loss: 0.6694 (0.6522)  acc1: 83.2031 (82.9545)  acc5: 96.0938 (96.1293)  time: 0.6346  data: 0.1988  max mem: 16096\n",
      "Test:  [ 20/196]  eta: 0:01:30  loss: 0.6738 (0.6733)  acc1: 82.0312 (82.4777)  acc5: 95.7031 (95.7961)  time: 0.3928  data: 0.0001  max mem: 16096\n",
      "Test:  [ 30/196]  eta: 0:01:18  loss: 0.6219 (0.6322)  acc1: 84.3750 (83.9718)  acc5: 95.7031 (96.0181)  time: 0.3859  data: 0.0001  max mem: 16096\n",
      "Test:  [ 40/196]  eta: 0:01:10  loss: 0.6801 (0.6750)  acc1: 81.6406 (82.5934)  acc5: 95.7031 (95.9604)  time: 0.3861  data: 0.0001  max mem: 16096\n",
      "Test:  [ 50/196]  eta: 0:01:04  loss: 0.6937 (0.6724)  acc1: 80.0781 (82.3529)  acc5: 96.8750 (96.0938)  time: 0.3834  data: 0.0001  max mem: 16096\n",
      "Test:  [ 60/196]  eta: 0:00:58  loss: 0.7149 (0.6849)  acc1: 80.0781 (81.9864)  acc5: 96.4844 (96.1066)  time: 0.3854  data: 0.0001  max mem: 16096\n",
      "Test:  [ 70/196]  eta: 0:00:53  loss: 0.6616 (0.6716)  acc1: 80.8594 (82.2843)  acc5: 96.4844 (96.1983)  time: 0.3859  data: 0.0001  max mem: 16096\n",
      "Test:  [ 80/196]  eta: 0:00:48  loss: 0.6510 (0.6968)  acc1: 81.2500 (81.7467)  acc5: 95.7031 (95.9201)  time: 0.3860  data: 0.0001  max mem: 16096\n",
      "Test:  [ 90/196]  eta: 0:00:44  loss: 0.9469 (0.7444)  acc1: 76.1719 (80.6834)  acc5: 92.5781 (95.4370)  time: 0.3868  data: 0.0001  max mem: 16096\n",
      "Test:  [100/196]  eta: 0:00:39  loss: 1.1594 (0.7964)  acc1: 70.7031 (79.5521)  acc5: 90.6250 (94.8755)  time: 0.3864  data: 0.0001  max mem: 16096\n",
      "Test:  [110/196]  eta: 0:00:35  loss: 1.1594 (0.8214)  acc1: 72.2656 (79.0365)  acc5: 91.4062 (94.6298)  time: 0.3836  data: 0.0001  max mem: 16096\n",
      "Test:  [120/196]  eta: 0:00:31  loss: 0.9820 (0.8389)  acc1: 76.1719 (78.7771)  acc5: 92.1875 (94.3634)  time: 0.3856  data: 0.0001  max mem: 16096\n",
      "Test:  [130/196]  eta: 0:00:26  loss: 1.0825 (0.8705)  acc1: 72.6562 (77.9610)  acc5: 91.4062 (94.0303)  time: 0.3866  data: 0.0001  max mem: 16096\n",
      "Test:  [140/196]  eta: 0:00:22  loss: 1.1088 (0.8889)  acc1: 72.2656 (77.6125)  acc5: 91.4062 (93.8137)  time: 0.3879  data: 0.0001  max mem: 16096\n",
      "Test:  [150/196]  eta: 0:00:18  loss: 1.1069 (0.9059)  acc1: 73.0469 (77.2998)  acc5: 91.0156 (93.5586)  time: 0.3914  data: 0.0002  max mem: 16096\n",
      "Test:  [160/196]  eta: 0:00:14  loss: 1.1360 (0.9197)  acc1: 73.0469 (77.0380)  acc5: 90.2344 (93.3472)  time: 0.3898  data: 0.0002  max mem: 16096\n",
      "Test:  [170/196]  eta: 0:00:10  loss: 1.2171 (0.9371)  acc1: 71.8750 (76.6265)  acc5: 89.8438 (93.1721)  time: 0.3845  data: 0.0002  max mem: 16096\n",
      "Test:  [180/196]  eta: 0:00:06  loss: 1.2493 (0.9527)  acc1: 68.7500 (76.2992)  acc5: 90.2344 (93.0205)  time: 0.3815  data: 0.0001  max mem: 16096\n",
      "Test:  [190/196]  eta: 0:00:02  loss: 1.0816 (0.9511)  acc1: 71.4844 (76.3089)  acc5: 92.1875 (93.0465)  time: 0.3736  data: 0.0001  max mem: 16096\n",
      "Test: Total time: 0:01:17\n",
      " * Acc@1 76.426 Acc@5 93.080\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    evaluate(model, criterion, data_loader_test, device=\"cuda\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After only 1 epoch of quantized fine tuning, top-1 improved from ~76.1 to 76.426. Train longer with lr anealing can improve accuracy futher"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
